import mujoco
import mujoco.viewer as viewer
import numpy as np
from os.path import dirname, join
import time

# 导入自定义函数
from os.path import dirname, join
import sys
sys.path.append(join(dirname(__file__), "src"))
sys.path.append(dirname(__file__))
from src import *


# mujoco加载模型
dual_arm_model= mujoco.MjModel.from_xml_path(join(dirname(__file__),'scene_dual_arm.xml')) 
dual_arm_data = mujoco.MjData(dual_arm_model)
mocap_id_dual_arm = dual_arm_model.body("target").mocapid[0]


pos_sensor_names = [
        "Motor_Shoulder_posR", 
        "Motor_BigArm1_posR",
        "Motor_SmallArm1_posR",
        "Motor_SmallArm2_posR",
        "Motor_Hand_posR",
        "Motor_Finger_posR",

        "Motor_Shoulder_posL", 
        "Motor_BigArm1_posL",
        "Motor_SmallArm1_posL",
        "Motor_SmallArm2_posL",
        "Motor_Hand_posL",
        "Motor_Finger_posL"
        ]


# head x轴指向右臂，z竖直向上
l = 0.3
roT_h = np.array([[0, 1, 0, 0],[0, 0, 1, 0],[1, 0, 0, -l],[0, 0, 0, 1]])
loT_h = np.array([[0, 1, 0, 0],[0, 0, -1, 0],[-1, 0, 0, -l],[0, 0, 0, 1]])

hT_e = np.eye(4)


box_pos = np.array([0, 0.4, -0.5])
box_quat = np.array([1,0,0,0])

with viewer.launch_passive(dual_arm_model,dual_arm_data) as viewer:
    # Initialize the camera view to that of the free camera.
    mujoco.mjv_defaultFreeCamera(dual_arm_model, viewer.cam)
    viewer.opt.frame = mujoco.mjtFrame.mjFRAME_SITE

    # 初始化位置
    loop_start  = time.time()
    while viewer.is_running():
        q_real_float = np.array([dual_arm_data.sensor(pos_sensor_name).data[0] for pos_sensor_name in pos_sensor_names])
        q_real_float_r = q_real_float[0:6]
        q_real_float_l = q_real_float[6:12]
        


        start = time.time()

        if time.time()-loop_start<3:
            box_width = 1
        else:
            box_width = 0.5

        anglebias = np.pi*np.clip(30, -30, 30)/180

        head_pos = np.array([0,0,0])
        c, s = np.cos(anglebias), np.sin(anglebias)
        head_rot = np.array([
            [1,  0,   0],
            [0,  c,  -s],
            [0,  s,   c]
        ])
        hT_e[:3, 3] = -head_rot.T@head_pos
        hT_e[:3, :3] = head_rot.T


        rd_pos, rd_quat, ld_pos, ld_quat = pick_box_r_l(box_width, box_pos, box_quat)

        roT_e = roT_h@hT_e
        roP_D, roR_D = eD2oD(rd_pos, rd_quat, roT_e)

        loT_e = loT_h@hT_e
        loP_D, loR_D = eD2oD(ld_pos, ld_quat, loT_e)



        rd_pos, rd_quat, ld_pos, ld_quat = pick_box_r_l(box_width, box_pos, box_quat)


        
        roT_e = roT_h@hT_e
        roP_D, roR_D = eD2oD(rd_pos, rd_quat, roT_e)

        
        loT_e = loT_h@hT_e
        loP_D, loR_D = eD2oD(ld_pos, ld_quat, loT_e)
        
        ctrl_r = ik_general(q_real_float_r, roP_D, roR_D)
        ctrl_l = ik_general(q_real_float_l, loP_D, loR_D)

        dual_arm_data.ctrl = np.hstack((motor_angle_clip(ctrl_r), motor_angle_clip(ctrl_l)))




        mujoco.mj_step(dual_arm_model, dual_arm_data)
        mujoco.mj_forward(dual_arm_model, dual_arm_data)

        viewer.sync()
        time.sleep(dual_arm_model.opt.timestep)